import os, sys
from typing import Dict

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 4)))
sys.path.insert(0, parent_path)

from base64 import b64encode
import traceback
import numpy as np
import socket, json, cv2, yaml, time
from typing import *
from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.PIL_draw import pic_text
from core.utils.visualize import get_color_map_list
from collections import deque

class RTMPProcessor(BaseProcessor):
    """直播流检测后处理方案
    """

    def _init_check(self):
        pass

    def _call(self, data:Dict) -> Dict:
        img_data = data[self.keys["in"]]
        det_result = data[self.keys["det"]]
        idx = 0
        for img_data, bboxes in zip(img_data, det_result["boxes"]):
            img = img_data.copy()
            draw_thickness = min(img.shape[:2]) // 320
            for box in bboxes:
                cls_id, score, x1,y1,x2,y2 = box
                cls_id = int(cls_id)
                img = cv2.rectangle(img,(int(x1), int(y1)),(int(x2),int(y2)),color=[0,0,255],thickness=draw_thickness)
            cv2.imwrite("temp_{}.jpg".format(idx),img)
            idx+=1
            

        return data
    
class RTMPTriggerProcessor(BaseProcessor):
    """直播流检测后处理方案
    """

    continue_threshold = 5
    continue_list = deque(maxlen=5)
    last_trigger_time = time.time()

    def _init_check(self):
        pass

    def _call(self, data:Dict) -> Dict:

        det_result = data[self.keys["det"]]
        need_trigger = False
        for bboxes in det_result["boxes"]:
            have_target = False
            for box in bboxes:
                cls_id, score, x1,y1,x2,y2 = box
                cls_id = int(cls_id)
                if cls_id in [1,2,3]:
                    have_target = True
                    break
            self.continue_list.append(have_target)
            conti = self.continue_list.count(1)
            if time.time() - self.last_trigger_time>=20:
                if conti>=self.continue_threshold:
                    need_trigger = True
                    self.last_trigger_time = time.time()
        data[self.keys["trigger_tag"]]=need_trigger
            

        return data
    
    